import os
import json
import requests
import threading
from concurrent.futures import ThreadPoolExecutor
import time
from urllib.parse import urlparse
from pathlib import Path

# 图片保存目录
image_dir = 'static/img/products'

# 线程锁，用于打印信息和写入统计信息
print_lock = threading.Lock()
stats_lock = threading.Lock()

# 统计信息
stats = {
    'total': 0,
    'success': 0,
    'failed': 0,
    'skipped': 0
}

def ensure_dir_exists(directory):
    """确保目录存在，如果不存在则创建"""
    if not os.path.exists(directory):
        os.makedirs(directory)
        print(f"已创建目录: {directory}")

def get_file_extension(url):
    """从URL中获取文件扩展名"""
    parsed_url = urlparse(url)
    path = parsed_url.path
    # 默认使用jpg，如果能从URL中提取到扩展名则使用提取的扩展名
    extension = os.path.splitext(path)[1]
    if extension and len(extension) > 1:
        return extension
    return ".jpg"  # 默认使用jpg扩展名

def download_image(skuid, image_url):
    """下载单个图片并保存"""
    try:
        # 构建图片保存路径
        extension = get_file_extension(image_url)
        image_filename = f"{skuid}_main{extension}"
        image_path = os.path.join(image_dir, image_filename)
        
        # 检查文件是否已存在，如果存在则跳过
        if os.path.exists(image_path):
            with stats_lock:
                stats['skipped'] += 1
            with print_lock:
                print(f"已跳过 {skuid}，文件已存在: {image_path}")
            return True
        
        # 下载图片
        response = requests.get(image_url, timeout=10)
        if response.status_code != 200:
            with print_lock:
                print(f"下载失败 {skuid}: HTTP 状态码 {response.status_code}")
            with stats_lock:
                stats['failed'] += 1
            return False
        
        # 保存图片
        with open(image_path, 'wb') as f:
            f.write(response.content)
        
        with stats_lock:
            stats['success'] += 1
        with print_lock:
            print(f"已下载 {skuid} 的图片: {image_path}")
        return True
    except Exception as e:
        with stats_lock:
            stats['failed'] += 1
        with print_lock:
            print(f"下载 {skuid} 时出错: {str(e)}")
        return False

def load_json_data(json_file):
    """加载JSON文件中的数据"""
    try:
        with open(json_file, 'r', encoding='utf-8') as f:
            data = json.load(f)
        with print_lock:
            print(f"已加载 {len(data)} 个商品图片URL")
        return data
    except Exception as e:
        with print_lock:
            print(f"加载JSON文件时出错: {str(e)}")
        return {}

def main():
    """主函数"""
    start_time = time.time()
    
    # 确保图片保存目录存在
    ensure_dir_exists(image_dir)
    
    # 加载JSON数据
    json_file = '主图.json'
    image_data = load_json_data(json_file)
    
    if not image_data:
        print("没有找到图片数据，请检查JSON文件")
        return
    
    with stats_lock:
        stats['total'] = len(image_data)
    
    # 创建下载任务列表
    tasks = [(skuid, url) for skuid, url in image_data.items()]
    
    # 使用线程池下载图片
    with ThreadPoolExecutor(max_workers=20) as executor:
        executor.map(lambda x: download_image(x[0], x[1]), tasks)
    
    # 统计结果
    end_time = time.time()
    duration = end_time - start_time
    
    print(f"\n下载统计:")
    print(f"总计: {stats['total']}")
    print(f"成功: {stats['success']}")
    print(f"失败: {stats['failed']}")
    print(f"跳过: {stats['skipped']}")
    print(f"总耗时: {duration:.2f} 秒")
    print(f"平均速度: {stats['success'] / duration:.2f} 张/秒")
    
    # 保存统计结果到文件
    with open('download_stats.txt', 'w', encoding='utf-8') as f:
        f.write(f"下载统计:\n")
        f.write(f"总计: {stats['total']}\n")
        f.write(f"成功: {stats['success']}\n")
        f.write(f"失败: {stats['failed']}\n")
        f.write(f"跳过: {stats['skipped']}\n")
        f.write(f"总耗时: {duration:.2f} 秒\n")
        f.write(f"平均速度: {stats['success'] / duration:.2f} 张/秒\n")
        f.write(f"完成时间: {time.strftime('%Y-%m-%d %H:%M:%S')}\n")

if __name__ == "__main__":
    main()
